# -*- coding: utf-8 -*-

from __future__ import print_function  # 兼容 Python 2 和 3 的 print 函数
import argparse
import os
import subprocess
import elfcheckSymboldiffdb
import time

RPM_PACKAGE_PATH = "/tmp/01_symbol_checker_rpm"
TAR_PACKAGE_PATH = "/tmp/02_symbol_checker_tar"


def unCompressRpm(RpmPath, unCompressPath):
    if not os.path.exists(unCompressPath):
        os.makedirs(unCompressPath)
    
    copy_cmd = "cp {} {}/".format(RpmPath, unCompressPath)
    os.system(copy_cmd)
    cmd = "cd {} && rpm2cpio {}/$(basename {}) | cpio -div --no-absolute-filenames".format(unCompressPath, unCompressPath, RpmPath)
    os.system(cmd)


def unCompressTar(TarPath, unCompressPath):
    """解压 tar.gz 或 zip 压缩包"""
    if not os.path.exists(unCompressPath):
        os.makedirs(unCompressPath)
        
    result = subprocess.check_output(['file', '-b', TarPath])
    file_type = result.decode('utf-8').strip()  # 将字节流解码为字符串
    print("File type: {}".format(file_type))
        
    # 根据文件类型判断并执行相应的解压命令
    if 'Zip archive data' in file_type:
        # 如果是 zip 文件
        cmd = "unzip -o {} -d {}".format(TarPath, unCompressPath)
        print("unCompressTar (zip): {}".format(cmd))
        os.system(cmd)
        
    elif 'gzip compressed data' in file_type:
        # 如果是 tar.gz 文件
        cmd = "tar -xzvf {} -C {}".format(TarPath, unCompressPath)
        print("unCompressTar (tar.gz): {}".format(cmd))
        os.system(cmd)
        
    else:
        print("Unsupported file type: {}".format(file_type))
   

def ScanFileInDir(ScanPath):
    filelist = []
    for root, dirs, files in os.walk(ScanPath):
        for file in files:
            filelist.append(os.path.join(root, file))
    return filelist 

def IsElfFormatFile(FilePath):
    process = subprocess.Popen(['file', FilePath], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()

    strResult = stdout.decode('utf-8') if hasattr(stdout, 'decode') else stdout  # 兼容 Python 2
    return "ELF 64-bit" in strResult

if __name__ == '__main__':
    start_time = time.time()
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('-P', type=str, help='支持单个RPM包')
    parser.add_argument('-D', type=str, help='支持单条指定路径')
    parser.add_argument('-C', type=str, help="支持单个压缩包")
    parser.add_argument('-E', type=str, help="支持单个ELF格式文件")
    param = parser.parse_args()
    
    RPM_Package = param.P
    ELF_Dir = param.D
    Compress_File = param.C
    ELF_File = param.E
    print("{} {} {} {}".format(RPM_Package, ELF_Dir, Compress_File, ELF_File))

    if RPM_Package is not None:
        if os.path.exists(RPM_Package):
            if os.path.exists(RPM_PACKAGE_PATH):
                os.system("rm -rf {}".format(RPM_PACKAGE_PATH))

            unCompressRpm(RPM_Package, RPM_PACKAGE_PATH)
            filelist = ScanFileInDir(RPM_PACKAGE_PATH)
            for item in filelist:
                if IsElfFormatFile(item):
                    print("扫描出来的二进制文件{}".format(item))
                    elfcheckSymboldiffdb.elf_check_symboldiffdb(item)
        else:
            print("{} RPM文件不存在".format(RPM_Package))

    if ELF_Dir is not None:
        if os.path.exists(ELF_Dir):
            filelist = ScanFileInDir(ELF_Dir)
            for item in filelist:
                if IsElfFormatFile(item):
                    elfcheckSymboldiffdb.elf_check_symboldiffdb(item)
        else:
            print("{} 路径不存在".format(ELF_Dir))
    
    if Compress_File is not None:
        if os.path.exists(Compress_File):
            if os.path.exists(TAR_PACKAGE_PATH):
                os.system("rm -rf {}".format(TAR_PACKAGE_PATH))

            unCompressTar(Compress_File, TAR_PACKAGE_PATH)
            filelist = ScanFileInDir(TAR_PACKAGE_PATH)
            for item in filelist:
                if IsElfFormatFile(item):
                    elfcheckSymboldiffdb.elf_check_symboldiffdb(item)
        else:
            print("{} 压缩包不存在".format(Compress_File))
    
    if ELF_File is not None:
        if os.path.exists(ELF_File):
            if IsElfFormatFile(ELF_File):
                elfcheckSymboldiffdb.elf_check_symboldiffdb(ELF_File)
        else:
            print("{} ELF文件不存在".format(ELF_File))
    
    end_time = time.time()
    execution_time = end_time - start_time
    elfcheckSymboldiffdb.close_db_connections()
    print("Execution time: {:.4f} seconds".format(execution_time))
